use std::sync::atomic::AtomicUsize;
use std::sync::mpsc::{Receiver, Sender};
use std::sync::Arc;
use std::thread::JoinHandle;
use std::collections::BTreeMap;

use super::exec::*;
use super::error::*;

pub trait MspcTaskAttr: Send + Sync + std::fmt::Debug + Clone + 'static {
    fn check_specific_name(name:&String) -> bool;
    fn schedule(sched:&MpscTaskScheduler<Self>) -> Result<(), <MpscTaskScheduler<Self> as TaskScheduler<Self>>::Error>;
}

pub struct MpscTaskScheduler<T:MspcTaskAttr> {
    handles:Vec<JoinHandle<()>>,
    task_counter:Arc<AtomicUsize>,
    exec_sender:Sender<Option<T>>,
    sched_recv:Receiver<Option<T>>,
    dyn_sched_sender:Vec<Sender<T>>,
    specific_sched_sender:BTreeMap<String, Sender<T>>,
    error_collector:ErrorCollector<ScheduleError>,
}

impl<T:MspcTaskAttr> MpscTaskScheduler<T> {
    pub fn new() -> Self {
        let (exec_sender, sched_recv) = std::sync::mpsc::channel::<Option<T>>();
        let error_collector = ErrorCollector::new();
        Self { 
            handles:Vec::new(), 
            task_counter:Arc::new(AtomicUsize::new(0)),
            exec_sender, sched_recv, error_collector,
            dyn_sched_sender:Vec::new(),
            specific_sched_sender:BTreeMap::new(),
        }
    }

    pub fn task_counter(&self) -> Arc<AtomicUsize> {
        self.task_counter.clone()
    }

    pub fn get_specific_sched_sender<Q>(&self, name:Q) -> Option<&Sender<T>> where Q:ToString{
        self.specific_sched_sender.get(&name.to_string())
    }

    pub fn get_dyn_sched_sender(&self) -> &[Sender<T>] {
        &self.dyn_sched_sender
    }

    pub fn get_sched_recv(&self) -> &Receiver<Option<T>> {
        &self.sched_recv
    }

    pub fn send_init_task(&self, task:T){
        self.task_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        self.exec_sender.send(Some(task)).expect("mspc scheduler init task error")
    }

    pub fn join(&mut self) {
        while !self.handles.is_empty() {
            self.handles.pop().unwrap().join().expect("join handle error");
        }
    }

    pub fn drop_sender(&mut self) {
        while let Some(sender) = self.dyn_sched_sender.pop() {
            drop(sender);
        }
        while let Some(sender) = self.specific_sched_sender.pop_first() {
            drop(sender);
        }
    }
}

impl<T:MspcTaskAttr> TaskScheduler<T> for MpscTaskScheduler<T> {
    type Receiver = Receiver<T>;
    type Sender = Sender<Option<T>>;
    type Error = ScheduleError;
    fn recv(recv:&Self::Receiver) -> Option<T> { recv.recv().ok() }
    fn send_with_count(sender:&Self::Sender, task:T, task_counter:&AtomicUsize) { 
        task_counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        sender.send(Some(task)).expect("scheduler send error"); 
    }
    fn terminal(sender:&Self::Sender) { sender.send(None).expect("scheduler send error"); }
    fn register_dyn_executor<TE:TaskExecutor<T>>(&mut self, executor:TE) {
        let (sched_sender, exec_recv) = std::sync::mpsc::channel::<T>();
        let task_counter = self.task_counter.clone();
        let error_collector = self.error_collector.clone();
        let exec_sender = self.exec_sender.clone();
        let handle = std::thread::spawn(move || {
            let run_result = executor.run::<Self>(exec_recv, exec_sender, task_counter.clone());
            error_collector.collect(run_result.map_err(|e|format!("{e:?}").into()))
        });
        self.dyn_sched_sender.push(sched_sender);
        self.handles.push(handle);
    }
    fn register_specific_executor<TE:TaskExecutor<T>>(&mut self, name:String, executor:TE) -> bool {
        if !T::check_specific_name(&name) { return false; }
        let (sched_sender, exec_recv) = std::sync::mpsc::channel::<T>();
        let task_counter = self.task_counter.clone();
        let error_collector = self.error_collector.clone();
        let exec_sender = self.exec_sender.clone();
        let handle = std::thread::spawn(move || {
            let run_result = executor.run::<Self>(exec_recv, exec_sender, task_counter.clone());
            error_collector.collect(run_result.map_err(|e|format!("{e:?}").into()))
        });
        self.specific_sched_sender.insert(name, sched_sender);
        self.handles.push(handle);
        true
    }

    fn schedule(&mut self) -> Result<(), Self::Error> {
        T::schedule(self)?;
        self.drop_sender();
        self.join();
        Ok(())
    }

    fn error_collector(&self) -> ErrorCollector<Self::Error> {
        self.error_collector.clone()
    }
}